from mskpp import vadd, Tensor, Chip
def my_mmad(gm_x, gm_y, gm_z):

    ub_x, ub_y, ub_z = Tensor("UB"), Tensor("UB"), Tensor("UB")
    ub_x.load(gm_x)
    ub_y.load(gm_y)
    out = vadd(ub_x, ub_y, ub_z)()
    return ub_z

if __name__ == '__main__':
    with Chip("Ascend910B4") as chip:
        chip.enable_trace()    # 使能算子模拟流水图的功能，生成trace.json文件
        chip.enable_metrics()   # 使能单指令及分PIPE的流水信息，生成Instruction_statistic.csv和Pipe_statistic.csv文件
        # 这里进入了对数据切分逻辑的处理，对一大块GM的数据，如何经过拆分成小数据分批次搬入，如何对
        # 内存进行分片多buffer搬运，都是属于tiling策略的范畴，这里模拟了单buffer情况，
        in_x = Tensor("GM", "FP16", [64, 64], format="ND")
        in_y = Tensor("GM", "FP16", [64, 64], format="ND")
        in_z = Tensor("GM", "FP16", [64, 64], format="ND")
        out_z = my_mmad(in_x, in_y, in_z)
        in_z.load(out_z)
